import heapq
from typing import List


# 堆 O(NlogN)


class Solution:
    def assignTasks(self, servers: List[int], tasks: List[int]) -> List[int]:
        working = []  # [结束时间,服务器权重,服务器ID]
        waiting = [(server, i) for i, server in enumerate(servers)]  # [服务器权重,服务器ID]
        heapq.heapify(waiting)

        ans = []
        for j, task in enumerate(tasks):
            now = j
            while working and working[0][0] <= j:
                _, server, i = heapq.heappop(working)
                heapq.heappush(waiting, (server, i))

            if not waiting:
                end, server, i = heapq.heappop(working)
                heapq.heappush(waiting, (server, i))
                now = end

            server, i = heapq.heappop(waiting)
            ans.append(i)
            heapq.heappush(working, (now + task, server, i))

        return ans


if __name__ == "__main__":
    print(Solution().assignTasks(servers=[3, 3, 2], tasks=[1, 2, 3, 2, 1, 2]))  # [2,2,0,2,1,2]
    print(Solution().assignTasks(servers=[5, 1, 4, 3, 2], tasks=[2, 1, 2, 4, 5, 2, 1]))  # [1,4,1,4,1,3,2]

    # 测试用例26/34
    # [6,9,5,4,10,5,0,8,4,2,11,9,3,7,1,4,0,4,1,8]
    print(Solution().assignTasks(servers=[31, 96, 73, 90, 15, 11, 1, 90, 72, 9, 30, 88],
                                 tasks=[87, 10, 3, 5, 76, 74, 38, 64, 16, 64, 93, 95, 60, 79, 54, 26, 30, 44, 64, 71]))
